"""
Point cloud registration algorithms and helpers, adapted to the new project layout.
"""

from __future__ import annotations

import copy
import random
import time
from dataclasses import dataclass
from typing import Optional

import numpy as np
import open3d as o3d

from config.settings import settings


@dataclass
class RegistrationResult:
    transformation: np.ndarray
    max_distance: float
    min_distance: float
    mean_distance: float


def compute_average_spacing(pcd: o3d.geometry.PointCloud, sample_points: int = 1000) -> float:
    if len(pcd.points) < 2:
        return 0.0

    if len(pcd.points) > sample_points:
        indices = random.sample(range(len(pcd.points)), sample_points)
        sampled_pcd = pcd.select_by_index(indices)
    else:
        sampled_pcd = pcd

    kdtree = o3d.geometry.KDTreeFlann(sampled_pcd)
    distances = []
    for i in range(len(sampled_pcd.points)):
        _, _, dist = kdtree.search_knn_vector_3d(sampled_pcd.points[i], 2)
        if len(dist) > 1:
            distances.append(dist[1])
    return float(np.mean(distances)) if distances else 0.0


def compute_dynamic_params(pcd: o3d.geometry.PointCloud, is_small_cloud: bool = False) -> dict[str, float]:
    bbox = pcd.get_axis_aligned_bounding_box()
    diagonal = np.linalg.norm(bbox.get_max_bound() - bbox.get_min_bound())
    avg_spacing = compute_average_spacing(pcd)

    if is_small_cloud:
        return {
            "diagonal": diagonal,
            "avg_spacing": avg_spacing,
            "ransac_max_correspondence": max(avg_spacing * 2.5, diagonal * 0.05),
            "fpfh_radius": max(avg_spacing * 5, diagonal * 0.1),
            "normal_radius": max(avg_spacing * 3, diagonal * 0.02),
            "ransac_iterations": min(400_000, len(pcd.points) * 100),
            "ransac_confidence": 0.999,
            "max_nn": 80,
        }

    return {
        "diagonal": diagonal,
        "avg_spacing": avg_spacing,
        "ransac_max_correspondence": max(avg_spacing * 3, diagonal * 0.005),
        "fpfh_radius": max(avg_spacing * 4, diagonal * 0.01),
        "normal_radius": max(avg_spacing * 2, diagonal * 0.005),
        "ransac_iterations": min(200_000, len(pcd.points) * 50),
        "ransac_confidence": 0.99,
        "max_nn": 50,
    }


def preprocess_point_cloud(
    pcd: o3d.geometry.PointCloud, is_small_cloud: bool = False, voxel_size: Optional[float] = None
) -> o3d.geometry.PointCloud:
    pcd = pcd.remove_duplicated_points()

    if not is_small_cloud:
        pcd, _ = pcd.remove_statistical_outlier(nb_neighbors=20, std_ratio=2.0)

    if voxel_size is None:
        bbox = pcd.get_axis_aligned_bounding_box()
        diagonal = np.linalg.norm(bbox.get_max_bound() - bbox.get_min_bound())
        voxel_size = diagonal * (0.003 if is_small_cloud else 0.005)

    if not is_small_cloud or len(pcd.points) > 10_000:
        pcd = pcd.voxel_down_sample(voxel_size)

    return pcd


def _compute_fpfh_feature(
    pcd: o3d.geometry.PointCloud,
    radius: float,
    max_nn: int,
) -> o3d.pipelines.registration.Feature:
    return o3d.pipelines.registration.compute_fpfh_feature(
        pcd,
        o3d.geometry.KDTreeSearchParamHybrid(radius=radius, max_nn=max_nn),
    )


def registration_point_clouds(
    scan_filename: str,
    bim_filename: str,
    is_small_cloud: Optional[bool] = None,
    max_attempts: Optional[int] = None,
    random_seed: Optional[int] = None,
) -> np.ndarray:
    cfg = settings.registration
    is_small = cfg.is_small_cloud if is_small_cloud is None else is_small_cloud
    attempts = cfg.max_attempts if max_attempts is None else max_attempts
    seed = cfg.random_seed if random_seed is None else random_seed

    np.random.seed(seed)
    random.seed(seed)

    start = time.time()
    best_result: Optional[RegistrationResult] = None
    best_distance = float("inf")

    point_cloud_original = o3d.io.read_point_cloud(scan_filename)
    if point_cloud_original.is_empty():
        raise ValueError("点云为空或加载失败！")

    bim_model = o3d.io.read_triangle_mesh(bim_filename)
    if bim_model.is_empty():
        raise ValueError("BIM模型为空或加载失败！")

    sampling_multiplier = (
        cfg.bim_sampling_multiplier_small if is_small else cfg.bim_sampling_multiplier_large
    )
    bim_point_cloud_original = bim_model.sample_points_uniformly(
        number_of_points=len(point_cloud_original.points) * sampling_multiplier
    )

    print(f"测量点云包含 {len(point_cloud_original.points)} 个点")
    print(f"BIM点云包含 {len(bim_point_cloud_original.points)} 个点")

    for attempt in range(attempts):
        if attempts > 1:
            print(f"尝试 {attempt + 1}/{attempts}")
            np.random.seed(seed)
            random.seed(seed)

        point_cloud = preprocess_point_cloud(copy.deepcopy(point_cloud_original), is_small)
        bim_point_cloud = preprocess_point_cloud(copy.deepcopy(bim_point_cloud_original), is_small)

        pc_params = compute_dynamic_params(point_cloud, is_small)
        bim_params = compute_dynamic_params(bim_point_cloud, is_small)

        point_cloud.estimate_normals(
            search_param=o3d.geometry.KDTreeSearchParamHybrid(
                radius=pc_params["normal_radius"],
                max_nn=pc_params["max_nn"],
            )
        )
        bim_point_cloud.estimate_normals(
            search_param=o3d.geometry.KDTreeSearchParamHybrid(
                radius=bim_params["normal_radius"],
                max_nn=bim_params["max_nn"],
            )
        )

        source_feature = _compute_fpfh_feature(point_cloud, pc_params["fpfh_radius"], pc_params["max_nn"])
        target_feature = _compute_fpfh_feature(bim_point_cloud, bim_params["fpfh_radius"], bim_params["max_nn"])

        print("执行RANSAC粗配准...")
        ransac_result = o3d.pipelines.registration.registration_ransac_based_on_feature_matching(
            source=point_cloud,
            target=bim_point_cloud,
            source_feature=source_feature,
            target_feature=target_feature,
            mutual_filter=True,
            max_correspondence_distance=pc_params["ransac_max_correspondence"],
            estimation_method=o3d.pipelines.registration.TransformationEstimationPointToPoint(False),
            ransac_n=4,
            checkers=[
                o3d.pipelines.registration.CorrespondenceCheckerBasedOnEdgeLength(0.9),
                o3d.pipelines.registration.CorrespondenceCheckerBasedOnDistance(
                    pc_params["ransac_max_correspondence"] * 1.2
                ),
            ],
            criteria=o3d.pipelines.registration.RANSACConvergenceCriteria(
                pc_params["ransac_iterations"],
                pc_params["ransac_confidence"],
            ),
        )

        print("RANSAC transformation matrix:")
        print(ransac_result.transformation)

        print("执行分级ICP精配准...")
        current_transformation = ransac_result.transformation

        if is_small:
            icp_stages = [
                {"name": "粗配准ICP", "max_distance": pc_params["diagonal"] * 0.1, "max_iteration": 2000},
                {"name": "中等精度ICP", "max_distance": pc_params["diagonal"] * 0.03, "max_iteration": 1000},
                {"name": "精配准ICP", "max_distance": pc_params["avg_spacing"] * 2, "max_iteration": 500},
            ]
        else:
            icp_stages = [
                {"name": "粗配准ICP", "max_distance": pc_params["diagonal"] * 0.08, "max_iteration": 1000},
                {"name": "精配准ICP", "max_distance": pc_params["avg_spacing"] * 2, "max_iteration": 200},
            ]

        for stage in icp_stages:
            print(f"执行{stage['name']}...")
            icp_result = o3d.pipelines.registration.registration_icp(
                source=point_cloud,
                target=bim_point_cloud,
                max_correspondence_distance=stage["max_distance"],
                init=current_transformation,
                estimation_method=o3d.pipelines.registration.TransformationEstimationPointToPoint(),
                criteria=o3d.pipelines.registration.ICPConvergenceCriteria(
                    max_iteration=stage["max_iteration"],
                    relative_fitness=1e-6,
                    relative_rmse=1e-6,
                ),
            )
            current_transformation = icp_result.transformation

        final_point_cloud = copy.deepcopy(point_cloud_original)
        final_point_cloud.transform(current_transformation)

        distances = np.asarray(final_point_cloud.compute_point_cloud_distance(bim_point_cloud_original))
        max_distance = float(np.max(distances))
        min_distance = float(np.min(distances))
        mean_distance = float(np.mean(distances))

        print("当前尝试的配准结果:")
        print(f"  - 最大距离: {max_distance:.6f}")
        print(f"  - 最小距离: {min_distance:.6f}")
        print(f"  - 平均距离: {mean_distance:.6f}")

        if mean_distance < best_distance:
            best_distance = mean_distance
            best_result = RegistrationResult(
                transformation=current_transformation,
                max_distance=max_distance,
                min_distance=min_distance,
                mean_distance=mean_distance,
            )

        if is_small and max_distance < 0.01:
            print("已达到良好配准精度，提前结束尝试")
            break

    if best_result is None:
        raise RuntimeError("配准失败，未获得有效结果")

    point_cloud_original.transform(best_result.transformation)

    total_time = time.time() - start
    print(f"配准总耗时：{total_time:.2f}秒 ({total_time / 60:.2f}分钟)")

    return best_result.transformation


